import asyncio
import logging
from typing import Dict, Any, Optional
from pathlib import Path
import yaml

from .data_storage import DataStorage
from .computation import ComputationManager
from .communication import CommunicationManager

class Worker:
    """Main worker class that coordinates storage, computation and communication.
    
    This class implements the main logic for worker nodes, managing task execution,
    data handling, and coordination with the master node.
    """
    
    def __init__(self, config_path: str):
        """Initialize the worker node.
        
        Args:
            config_path: Path to the configuration file
        """
        # Load configuration
        with open(config_path, 'r') as f:
            self.config = yaml.safe_load(f)

        # Initialize components
        # Find worker configuration based on hostname
        import socket
        hostname = socket.gethostname()
        worker_config = next(
            (w for w in self.config['system']['workers']
             if w['host'] == hostname),
            None
        )
        if not worker_config:
            raise RuntimeError(f"No worker configuration found for host {hostname}")
            
        self.worker_id = worker_config['worker_id']
        self.master_url = f"http://{self.config['system']['master']['host']}:{self.config['system']['master']['port']}"
        
        # Setup logging
        self._setup_logging()
        
        # Create cache directory
        cache_dir = Path(self.config.get('cache_dir', 'cache'))
        cache_dir.mkdir(parents=True, exist_ok=True)
        
        # Initialize components
        self.data_storage = DataStorage(
            cache_dir=str(cache_dir),
            max_cache_size=self.config['memory']['cache_size']
        )
        
        self.computation = ComputationManager(
            data_storage=self.data_storage,
            device=f"cuda:{self.config['system'].get('gpu_id', 0)}"
        )
        
        self.communication = CommunicationManager(
            worker_id=self.worker_id,
            master_url=self.master_url,
            data_storage=self.data_storage,
            config=self.config
        )
        
        self.running = False
        self.current_task: Optional[Dict[str, Any]] = None
    
    def _setup_logging(self) -> None:
        """Setup logging configuration."""
        logging.basicConfig(
            level=self.config['logging']['level'],
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            filename=self.config['logging']['file']
        )
        self.logger = logging.getLogger(f"Worker-{self.worker_id}")
    
    async def start(self) -> None:
        """Start the worker node."""
        self.logger.info("Starting worker node...")
        self.running = True
        
        try:
            # Start communication manager
            await self.communication.start()
            
            # Main task loop
            while self.running:
                await self._process_next_task()
                await self._report_status()
                await asyncio.sleep(1)  # Prevent tight loop
                
        except Exception as e:
            self.logger.error(f"Error in worker main loop: {str(e)}")
            raise
        finally:
            await self.stop()
    
    async def stop(self) -> None:
        """Stop the worker node."""
        self.logger.info("Stopping worker node...")
        self.running = False
        await self.communication.stop()
    
    async def _process_next_task(self) -> None:
        """Process the next available task."""
        if self.current_task is not None:
            return  # Still processing current task
            
        # Get next task
        task = await self.communication.get_task()
        if task is None:
            return
            
        self.current_task = task
        self.logger.info(f"Received task: {task['task_id']}")
        
        try:
            # Fetch required input data
            if not await self._fetch_task_inputs(task):
                raise RuntimeError("Failed to fetch input data")
            
            # Execute task
            result = await self._execute_task(task)
            
            # Update task status
            await self.communication.update_task_status(
                task_id=task['task_id'],
                status='completed',
                result=result
            )
            
        except Exception as e:
            self.logger.error(f"Error processing task {task['task_id']}: {str(e)}")
            await self.communication.update_task_status(
                task_id=task['task_id'],
                status='failed',
                result={'error': str(e)}
            )
        
        finally:
            self.current_task = None
    
    async def _fetch_task_inputs(self, task: Dict[str, Any]) -> bool:
        """Fetch all required input data for a task.
        
        Args:
            task: Task information
            
        Returns:
            bool: True if all inputs were fetched successfully
        """
        for input_info in task.get('inputs', []):
            data_id = input_info['data_id']
            source_worker = input_info['location']
            
            if not await self.communication.fetch_data(data_id, source_worker):
                self.logger.error(f"Failed to fetch input {data_id} from {source_worker}")
                return False
        
        return True
    
    async def _execute_task(self, task: Dict[str, Any]) -> Dict[str, Any]:
        """Execute a computation task.
        
        Args:
            task: Task information
            
        Returns:
            Dict[str, Any]: Task results
        """
        # Execute the computation
        output_id = self.computation.execute_layer(
            layer_id=task['layer_id'],
            input_ids=[input_info['data_id'] for input_info in task['inputs']],
            layer_config=task['config']
        )
        
        if output_id is None:
            raise RuntimeError("Computation failed")
        
        # Get memory usage statistics
        memory_stats = self.computation.get_memory_usage()
        
        return {
            'output_id': output_id,
            'memory_usage': memory_stats['allocated'],
            'output_shape': self.computation.get_output_shape(task['layer_id'])
        }
    
    async def _report_status(self) -> None:
        """Report worker status to master node."""
        metrics = {
            'memory_usage': self.computation.get_memory_usage(),
            'cache_stats': self.data_storage.get_cache_stats(),
            'current_task': self.current_task['task_id'] if self.current_task else None
        }
        
        await self.communication.report_metrics(metrics)

def main():
    """Main entry point for the worker process."""
    import argparse
    
    parser = argparse.ArgumentParser(description='Start a DeepSlicing worker node')
    parser.add_argument('--config', type=str, required=True,
                      help='Path to configuration file')
    
    args = parser.parse_args()
    
    # Create and start worker
    worker = Worker(args.config)
    
    try:
        asyncio.run(worker.start())
    except KeyboardInterrupt:
        print("\nShutting down worker...")
    except Exception as e:
        print(f"Error: {str(e)}")
        raise

if __name__ == '__main__':
    main()